// utils/chessNet.js
const tf = require('./tf.js');

class ChessNet {
  constructor() {
    // 初始化模型时明确指定输入形状
    this.model = tf.sequential();

    // 输入层：14个特征平面(6种棋子×2方+2个历史平面)
    this.model.add(tf.layers.conv2d({
      filters: 256,
      kernelSize: 3,
      padding: 'same',
      inputShape: [10, 9, 14], // 明确指定输入维度
      activation: 'relu'
    }));

    // 残差块×20
    for (let i = 0; i < 20; i++) {
      this.model.add(this.residualBlock());
    }

    // 策略头
    const policyHead = tf.sequential();
    policyHead.add(tf.layers.conv2d({
      filters: 2,
      kernelSize: 1,
      padding: 'same',
      activation: 'relu',
      inputShape: [10, 9, 256] // 动态适应输入;
    }));
    policyHead.add(tf.layers.flatten());
    policyHead.add(tf.layers.dense({
      units: 2086,
      activation: 'softmax' // 中国象棋平均约2086种合法走法
    }));

    // 价值头
    const valueHead = tf.sequential();
    valueHead.add(tf.layers.conv2d({
      filters: 1,
      kernelSize: 1,
      padding: 'same',
      activation: 'relu',
      inputShape: [10, 9, 256] // 动态适应输入;
    }));
    valueHead.add(tf.layers.flatten());
    valueHead.add(tf.layers.dense({
      units: 256,
      activation: 'relu'
    }));
    valueHead.add(tf.layers.dense({
      units: 1,
      activation: 'tanh'
    }));

    // 在残差块后添加形状转换层
    this.model.add(tf.layers.conv2d({
      filters: 256,
      kernelSize: 1,
      padding: 'same',
      activation: 'relu'
    }));
  }

  residualBlock() {
    // 创建输入层
    const input = tf.input({
      shape: [null, null, 256]
    });

    // 第一层卷积
    const conv1 = tf.layers.conv2d({
      filters: 256,
      kernelSize: 3,
      padding: 'same',
      activation: 'relu'
    }).apply(input);

    // 第二层卷积（无激活）
    const conv2 = tf.layers.conv2d({
      filters: 256,
      kernelSize: 3,
      padding: 'same'
    }).apply(conv1);

    // 残差连接
    const sum = tf.layers.add().apply([input, conv2]);

    // 最终激活
    const output = tf.layers.activation({
      activation: 'relu'
    }).apply(sum);

    return tf.model({
      inputs: input,
      outputs: output
    });
  }

  predict(input) {
    return tf.tidy(() => {
      const [policy, value] = this.model.predict(input);
      return {
        policy,
        value
      };
    });
  }
}

module.exports = ChessNet;